// Copyright 2016 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/cert/sth_distributor.h"

#include <map>
#include <string>

#include "base/test/histogram_tester.h"
#include "crypto/sha2.h"
#include "net/cert/signed_tree_head.h"
#include "net/cert/sth_observer.h"
#include "net/test/ct_test_util.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace net {

namespace ct {

    namespace {

        // An STHObserver implementation that simply stores all
        // observed STHs, keyed by log ID.
        class StoringSTHObserver : public STHObserver {
        public:
            void NewSTHObserved(const SignedTreeHead& sth) override
            {
                sths[sth.log_id] = sth;
            }

            std::map<std::string, SignedTreeHead> sths;
        };

        class STHDistributorTest : public ::testing::Test {
        public:
            STHDistributorTest() { }

            void SetUp() override
            {
                ASSERT_TRUE(GetSampleSignedTreeHead(&sample_sth_));
                sample_sth_.log_id = GetTestPublicKeyId();
            }

        protected:
            STHDistributor distributor_;
            SignedTreeHead sample_sth_;
        };

        // Test that when a new observer is registered, the STHDistributor notifies it
        // of all the observed STHs it received so far.
        // This test makes sure that all observed STHs are reported to the observer.
        TEST_F(STHDistributorTest, NotifiesOfExistingSTHs)
        {
            // Create an STH that differs from the |sample_sth_| by belonging to a
            // different log.
            const std::string other_log = "another log";
            SignedTreeHead second_sth(sample_sth_);
            second_sth.log_id = other_log;

            // Notify |distributor_| of both STHs.
            distributor_.NewSTHObserved(sample_sth_);
            distributor_.NewSTHObserved(second_sth);

            StoringSTHObserver observer;
            distributor_.RegisterObserver(&observer);

            // Check that two STHs from different logs received prior to observer
            // registration were reported to the observer once registered.
            EXPECT_EQ(2u, observer.sths.size());
            EXPECT_EQ(1u, observer.sths.count(other_log));
            distributor_.UnregisterObserver(&observer);
        }

        // Test that histograms are properly recorded for the STH age when an STH
        // from Google's Pilot log is observed.
        TEST_F(STHDistributorTest, LogsUMAForPilotSTH)
        {
            const char kPilotSTHAgeHistogram[] = "Net.CertificateTransparency.PilotSTHAge";
            base::HistogramTester histograms;
            histograms.ExpectTotalCount(kPilotSTHAgeHistogram, 0);

            const uint8_t kPilotLogID[] = {
                0xa4, 0xb9, 0x09, 0x90, 0xb4, 0x18, 0x58, 0x14, 0x87, 0xbb, 0x13,
                0xa2, 0xcc, 0x67, 0x70, 0x0a, 0x3c, 0x35, 0x98, 0x04, 0xf9, 0x1b,
                0xdf, 0xb8, 0xe3, 0x77, 0xcd, 0x0e, 0xc8, 0x0d, 0xdc, 0x10
            };
            sample_sth_.log_id = std::string(reinterpret_cast<const char*>(kPilotLogID),
                crypto::kSHA256Length);

            distributor_.NewSTHObserved(sample_sth_);
            histograms.ExpectTotalCount(kPilotSTHAgeHistogram, 1);
        }

        // Test that the STHDistributor updates, rather than accumulates, STHs
        // coming from the same log.
        // This is tested by notifying the STHDistributor of an STH, modifying that
        // STH, notifying the STHDistributor of the modified STH, then registering
        // an observer which should get notified only once, with the modified STH.
        TEST_F(STHDistributorTest, UpdatesObservedSTHData)
        {
            // Observe an initial STH
            StoringSTHObserver observer;
            distributor_.RegisterObserver(&observer);

            distributor_.NewSTHObserved(sample_sth_);

            EXPECT_EQ(1u, observer.sths.size());
            EXPECT_EQ(sample_sth_, observer.sths[GetTestPublicKeyId()]);

            // Observe a new STH. "new" simply means that it is a more recently observed
            // SignedTreeHead for the given log ID, not necessarily that it's newer
            // chronologically (the timestamp) or the log state (the tree size).
            // To make sure the more recently observed SignedTreeHead is returned, just
            // modify some fields.
            SignedTreeHead new_sth = sample_sth_;
            new_sth.tree_size++;
            new_sth.timestamp -= base::TimeDelta::FromSeconds(3);

            distributor_.NewSTHObserved(new_sth);
            // The STH should have been broadcast to existing observers.
            EXPECT_EQ(1u, observer.sths.size());
            EXPECT_NE(sample_sth_, observer.sths[GetTestPublicKeyId()]);
            EXPECT_EQ(new_sth, observer.sths[GetTestPublicKeyId()]);

            // Registering a new observer should only receive the most recently observed
            // STH.
            StoringSTHObserver new_observer;
            distributor_.RegisterObserver(&new_observer);
            EXPECT_EQ(1u, new_observer.sths.size());
            EXPECT_NE(sample_sth_, new_observer.sths[GetTestPublicKeyId()]);
            EXPECT_EQ(new_sth, new_observer.sths[GetTestPublicKeyId()]);

            distributor_.UnregisterObserver(&new_observer);
            distributor_.UnregisterObserver(&observer);
        }

    } // namespace

} // namespace ct

} // namespace net
